/* Janus Chat - Version 1.0
 * Authors: Matthew Perry
 *          Michael Freid
 *          Marc Latou
 *          Richard Lee
 *
 * The Janus Chat Client is a highly secured
 * peer-to-peer chat client designed to facilitate
 * the transfer of sensitive data or as a communication
 * medium for the extra paranoid
 *
 * This file contains the implementation for a Message class. 
 * This class will serve as the main packet encapsulation
 * for janus' communications */

#include "message.h"
#include "exceptions.h"

/* The message constructor:
 * This constructor takes one parameter,
 * a type. This constructor is used to construct
 * messages that have some extra fields such as
 * decrypted text messages which have a user
 * moniker field. The message construts the body
 * dynamically in this case from fields added by the
 * program. */
Message::Message( unsigned int type )
{
    //This constructor is only valid for decrypted text messages
    if( type != DecryptStandard && type < ServerInfo )
        throw new MessageException( "Improper message type" );

    //Initialize member variables
    myType = type;
    myFieldSize = 0;
}

/* The message constructor:
 * This constructor takes two parameters
 * a type and a body of text. This constructor is used
 * to construct messages whos data is handled elsewhere
 * such as key exchange messages. The entire body can be
 * supplied because the serialization of data occurs in
 * another area. */
Message::Message( unsigned int type, const string& body )
{
    //Size check on the body
    if( body.size() > BODY_SIZE )
        throw new MessageException( "Body is too large" );

    //Initialize member variables
    myType = type;
    myFieldSize = 0;
    myBody = body;
    myBodyLen = body.size();
}

/* The message constructor:
 * This constructor takes a character string and
 * a length and converts it into a message. The
 * function assumes the data is in the correct format
 * and extracts the appropriate fields. */
Message::Message( char* data, unsigned int length )
{
    //The length of the data must be a certain size
    if( length < MSG_SIZE )
        throw new MessageException( "Char array is too small" );

    myType = 0;
    /* Loop to reinitialize myType from char encoding */
    for( int i = 0; i < 4; ++i ) 
        myType = (myType << 8) | (unsigned char)data[i];

    myBodyLen = 0;
    // Loop to reinitialize myBodyLen from char encoding
    for( int i = 4; i < 8; ++i )
        myBodyLen = (myBodyLen << 8) | (unsigned char)data[i];
    
    /* Create new strings on the stack and construct
     * them from different offsets of the data */
    string body( data + 8, myBodyLen );
    myBody = body;

    // If this is a dynamic body message, restore fields
    if( myType == DecryptStandard || myType >= ServerInfo )
        fieldInit();
}

/* The encoder helper:
 * this function takes a type and a body and
 * constructs a character string encoding of that
 * data so that it can be sent easily over the network */
char* Message::encoder( unsigned int type, const string &body )
{
    char* msg = new char[ MSG_SIZE ];

    /* Loop to convert the myType integer into the first
     * four bytes of the outgoing c_string */
    for( int i = 3; i >= 0; --i )
        msg[i] = char( type >> ( ( 3 - i ) * 8 ) );

    /* Loop to convert the myBodyLen integer into the second
     * four bytes of the outgoing c_string */
    for( int i = 7; i >= 4; --i )
        msg[i] = char( body.size() >> ( ( 7 - i ) * 8 ) );

    /* Loop to convert the text body into the last portion
     * of the message */
    for( unsigned int i = 0; i <= body.size(); ++i )
        msg[i + 8] = body[i];

    /* Loop to pad the remaining bytes of the outgoing message
     * to keep a consistant MSG_SIZE bytes */
    for( unsigned int i = 0; i < ( (MSG_SIZE - 8) - body.size() ); ++i )
        msg[i + 8 + body.size()] = (char)0;

    return msg;
}

/* The bodyInit helper:
 * This is used by dynamic body messages. It
 * creates a body string from various fields. */
void Message::bodyInit()
{
    //This function is only valid for decrypted text messages
    if( myType != DecryptStandard && myType < ServerInfo )
        return;

    myBody.resize( 8 );
    //There must be exactly two fields
    if( myFields.size() != 2 )
        throw new MessageException( "Incorrect number of fields" );

    //Loop to convert field's size to and integer
    for( int i = 3; i >= 0; --i )
        myBody[i] = char( myFields[0].size() >> ( ( 3 - i ) * 8 ) );

    //Loop to convert field's size to and integer
    for( int i = 7; i >= 4; --i )
        myBody[i] = char( myFields[1].size() >> ( ( 7 - i ) * 8 ) );
 
    //Loop over the fields vector adding data of fields
    for( unsigned int i = 0; i < myFields.size(); ++i )
        myBody += myFields[i];

    myBodyLen = myBody.size();
}

/* The fieldInit helper:
 * This is used by dynamic body messages. It
 * creates fields from a string body */
void Message::fieldInit()
{
    //This function is only valid for decrypted text messages
    if( myType != DecryptStandard && myType < ServerInfo )
        return;

    int fieldLen1 = 0;
    /* Loop to reinitialize fieldlen1 from char encoding */
    for( int i = 0; i < 4; ++i )
        fieldLen1 = (fieldLen1 << 8) | (unsigned char)myBody[i];

    int fieldLen2 = 0;
    /* Loop to reinitialize fieldlen2 from char encoding */
    for( int i = 4; i < 8; ++i )
        fieldLen2 = (fieldLen2 << 8) | (unsigned char)myBody[i];

    string field1;
    /* Retrieve field1 data */
    for( int i = 0; i < fieldLen1; ++i )
       field1 += myBody[i + 8];
    myFields.push_back( field1 );

    string field2;
    /* Retrieve field2 data */
    for( int i = 0; i < fieldLen2; ++i )
        field2 += myBody[i + 8 + fieldLen1];
    myFields.push_back( field2 );
}

/* The message toChar function
 * this message encodes the member variables into
 * a single character string which can be
 * easily transmitted. */
char* Message::toChar()
{
    /* Create body from fields */
    bodyInit();
    
    // Encode data into char*
    return encoder( myType, myBody );
}

/* The message aesEncrypt function
 * this encrypts the body of the message
 * and encodes this with the type into
 * a single character string which can be
 * easily transmitted. */
char* Message::aesEncrypt( const string& key )
{
    /* Create body from fields */
    bodyInit();

    /* Generate the encrypted form of the body */
    string cipher = jAes::encrypt( key, myBody );

    // Encode data into char*
    return encoder( myType - 1, cipher );
}

/* The message aesDecrypt function
 * this decrypts the body of the message
 * and returns this with the type into
 * a message object. */
Message Message::aesDecrypt( const string& key ) const
{
    /* Decrypt body */
    string decryptedBody = jAes::decrypt( key, myBody );

    /* Create the message object corresponding to the decrypted message */
    Message decryptedMessage( myType + 1, decryptedBody );
    decryptedMessage.fieldInit();

    return decryptedMessage;
}

/* The message addField function:
 * This function adds a a field to the
 * message's field list for messages
 * that have a dynamic body */
void Message::addField( const string& str )
{
    //Calculate how long the body of a message with the current
    //fields would be if it were generated right now
    unsigned int proposedBodyLen = myFieldSize + ( 4 * myFields.size() ) + 4 + str.size();

    //This new field can only be added if the new proposed body
    //Length is under the maximum size of bodys
    if( proposedBodyLen > BODY_SIZE )
        throw new MessageException( "Field makes the body too long" );

    //If size isnt a problem, add the string as a field
    myFields.push_back( str );

    //Change fieldsize
    myFieldSize += str.size();
}

/* Static SessionKeyMessage function:
 * Static function which creates an RSA encrypted
 * message containing the AES session key for the
 * chat */
Message Message::SessionKeyMessage( const string& key, jRsa::publicKey pk )
{
    string cipherKey = jRsa::encrypt( key, pk );
    Message retMsg( Message::SessionKeyExchange, cipherKey );
    return retMsg;
}

/* Static SessionKeyRestore function:
 * Static function which decrypts the Session Key
 * exchange message that it is given and returns
 * the AES session key */
string Message::SessionKeyRestore( const Message& msg, jRsa::privateKey pk )
{
    if( msg.getType() != Message::SessionKeyExchange )
        throw new MessageException( "Not a Session Key Exchange Message" );

    string key = jRsa::decrypt( msg.myBody, pk );
    return key;
}
